from fastmcp import Client as FastMCPClient 
from pydantic import BaseModel
from typing import Optional, List, Dict, Any
import asyncio
from ConfigManager import ConfigManager
from JsonUtil import JsonUtil
from LLMProcess import LLMProcess
import json
import Logger_config

# 设置日志配置
logger = Logger_config.setup_logger()

class PlanningAgent:
    def __init__(self, client: FastMCPClient):
        self.client = client

        self.cm = ConfigManager()
        self.config = self.cm.get_merged_config()
    
    async def get_tools_description(self) -> str:
        """获取可用工具的描述信息"""
        tools = await self.client.list_tools()
        descriptions = []
        for tool in tools:
            desc = f"- {tool.name}: {tool.description}\n  Input schema: {tool.inputSchema}\n "
            descriptions.append(desc)
        return "\n".join(descriptions)
    
    async def plan_with_llm(self, goal: str) -> Dict[str, Any]:
        model_base_url = self.config['model']['base_url']  # 模型的基础URL
        model_api_key = self.config['model']['api_key']  # 模型的API密钥
        model_name = self.config['model']['model_name']  # 模型名称
        model_max_tokens = self.config['model']['max_tokens']  # 设置模型生成的最大token数，控制生成内容的长度
        model_temperature = self.config['model']['temperature']  # 设置模型的温度参数，控制生成内容的随机性，值越高生成的内容越随机
        model_stream = self.config['model']['stream']  # 设置是否以流式方式获取模型生成的内容，True表示逐步返回生成结果，False表示一次性返回

        """使用LLM生成执行计划"""
        tools_info = await self.get_tools_description()
        
        prompt = f"""你是一个智能规划助手，可以根据用户目标选择合适的工具并生成执行计划。

        可用工具:
        {tools_info}

        用户目标: {goal}

        请按照以下步骤操作:
        1. 分析目标并确定是否需要使用工具
        2. 如果不需要工具，直接返回最终答案
        3. 如果需要工具，选择最合适的工具并生成符合输入模式的参数
        4. 返回JSON格式的响应，包含以下字段:
        - "thoughts": 你的思考过程
        - "tool_required": 是否需要工具 (true/false)
        - "tool_name": 工具名称 (如果需要工具)
        - "tool_input": 工具输入参数 (如果需要工具)

        示例响应(需要工具):
        {{
            "thoughts": "用户想要查询活跃用户，可以使用query_database工具",
            "tool_required": true,
            "tool_name": "query_database",
            "tool_input": {{
                "query": "SELECT * FROM users WHERE active = true",
                "params": null
            }}
        }}

        示例响应(不需要工具):
        {{
            "thoughts": "用户只是打招呼，不需要使用任何工具",
            "tool_required": false,
            "response": "你好！我是你的AI助手。"
        }}

        现在请为以下目标生成响应:
        目标: {goal}
        """

        prompts=[
            {"role": "system", "content": "你是一个智能规划助手，可以根据用户目标选择合适的工具并生成执行计划。"},
            {"role": "user", "content": prompt}
        ]
        llm_result = LLMProcess.getLLMTextResponse(model_name,model_api_key,model_base_url,model_max_tokens,model_temperature,model_stream,prompts)
        result = JsonUtil.reshape(llm_result)#重新整理JSON格式
        
        try:
            return json.loads(result)
        except json.JSONDecodeError:
            logger.error(f"Invalid JSON response: {result}")
    
    async def execute_plan(self, plan: Dict[str, Any]) -> str:
        """执行LLM生成的计划并处理结果"""
        if not plan.get("tool_required", False):
            return plan.get("response", "No action required.")
        
        tool_name = plan["tool_name"]
        tool_input = plan["tool_input"]
        
        # 调用工具获取原始结果
        logger.info(f"Executing tool: {tool_name} with input: {tool_input}")
        try:
            raw_result = await self.client.call_tool(tool_name, tool_input)
        except Exception as e:
            logger.error(f"Tool execution failed: {str(e)}")
            return f"工具执行失败: {str(e)}"
        
        # 构造结果处理prompt
        process_prompt = f"""请根据以下信息生成友好的最终回答：
        用户原始目标: {plan.get('original_goal', '')}    
        模型thoughts: {plan.get('thoughts', '')}
        使用工具名称：{tool_name}
        工具输入参数: {tool_input}
        工具返回结果: {raw_result}
        
        要求：
        1. 根据工具返回结果、工具输入参数以及用户原始目标，生成一个友好的最终回答
        """
        model_base_url = self.config['model']['base_url']  # 模型的基础URL
        model_api_key = self.config['model']['api_key']  # 模型的API密钥
        model_name = self.config['model']['model_name']  # 模型名称
        model_max_tokens = self.config['model']['max_tokens']  # 设置模型生成的最大token数，控制生成内容的长度
        model_temperature = self.config['model']['temperature']  # 设置模型的温度参数，控制生成内容的随机性，值越高生成的内容越随机
        model_stream = self.config['model']['stream']  # 设置是否以流式方式获取模型生成的内容，True表示逐步返回生成结果，False表示一次性返回
        prompts=[
                {"role": "system", "content": "你是一个结果处理专家，可以将原始数据转换为友好回答"},
                {"role": "user", "content": process_prompt}
            ]
        processed_response = LLMProcess.getLLMTextResponse(model_name,model_api_key,model_base_url,model_max_tokens,model_temperature,model_stream,prompts)
        
        return processed_response
    
    async def plan_and_execute(self, goal: str) -> str:
        """完整的规划-执行流程"""
        # 第一步：使用LLM生成计划
        plan = await self.plan_with_llm(goal)
        plan["original_goal"] = goal  # 保留原始目标用于后续处理
        print(f"Generated plan: {json.dumps(plan, indent=2,ensure_ascii=False)}")
        
        # 第二步：执行计划
        if plan.get("tool_required", False):
            return await self.execute_plan(plan)
        else:
            return plan.get("response", "No action required.")

async def main():
    # 创建客户端
    async with FastMCPClient("http://localhost:8000/mcp") as client:
    
        # 创建规划代理
        agent = PlanningAgent(client)
        
        # 示例1: 数据库查询
        # result1 = await agent.plan_and_execute("查询活跃用户")
        # print("\nResult 1:", result1)
        
        # # 示例2: 数学计算
        # result2 = await agent.plan_and_execute("计算5.5加3.2")
        # print("\nResult 2:", result2)
        
        # # 示例3: 简单问候
        # result3 = await agent.plan_and_execute("你好")
        # print("\nResult 3:", result3)
        
        # # 示例4: 复杂查询
        # result4 = await agent.plan_and_execute("找出所有年龄大于30岁的用户")
        # print("\nResult 4:", result4)

        # 示例5: 查询日期
        result5 = await agent.plan_and_execute("明天天气怎样？你需要先使用工具获取明确的日期信息，然后再根据日期获取天气信息。")
        print("\nResult 4:", result5)

if __name__ == "__main__":
    asyncio.run(main())